import streamlit as st
import torch
from PIL import Image
from torchvision import transforms as T
from torchvision.models import resnet18
import torch.nn as nn
import os

# 设置页面
st.set_page_config(
    page_title="汽车品牌识别器",
    page_icon="🚗",
    layout="centered"
)

# 应用标题
st.title("🚗 汽车品牌识别器")
st.markdown("上传汽车图片，识别其品牌")


# 加载类别名称
@st.cache_data
def load_classnames():
    try:
        with open("./classname.txt", 'r', encoding='utf-8') as f:
            return {i: line.strip() for i, line in enumerate(f.readlines())}
    except:
        return {i: f"品牌{i}" for i in range(50)}


# 创建模型
@st.cache_resource
def load_model():
    device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')

    # 创建模型结构
    model = resnet18(weights=None)
    model.fc = nn.Linear(model.fc.in_features, 50)

    # 加载权重
    try:
        checkpoint = torch.load('D:\华清远见\shanghai_25071\w4机器视觉处理模型\项目\汽车分类\models\BEST_MODEL.pth', map_location=device)
        if isinstance(checkpoint, dict) and 'model_state_dict' in checkpoint:
            model.load_state_dict(checkpoint['model_state_dict'])
        else:
            model.load_state_dict(checkpoint)
        model.eval()
        model.to(device)
        return model, device
    except Exception as e:
        st.error(f"模型加载失败: {e}")
        return None, device


# 图像预处理
def preprocess_image(image):
    transform = T.Compose([
        T.Resize((128, 128)),
        T.ToTensor(),
        T.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225])
    ])
    return transform(image).unsqueeze(0)


# 预测函数
def predict(model, device, image_tensor):
    with torch.no_grad():
        image_tensor = image_tensor.to(device)
        outputs = model(image_tensor)
        probabilities = torch.softmax(outputs, dim=1)
        return probabilities


# 主应用
from 做标签映射表 import load_classname_classes
def main():
    # 加载模型和类别
    with st.spinner("加载模型中..."):
        model, device = load_model()
        classnames = load_classname_classes('D:\华清远见\shanghai_25071\w4机器视觉处理模型\项目\汽车分类\classname.txt')

    if model is None:
        st.error("无法加载模型，请检查模型文件是否存在")
        return

    # 图片上传
    uploaded_file = st.file_uploader(
        "选择汽车图片",
        type=['jpg', 'jpeg', 'png'],
        help="支持 JPG, JPEG, PNG 格式"
    )

    if uploaded_file is not None:
        # 显示图片
        image = Image.open(uploaded_file).convert('RGB')
        st.image(image, caption="上传的图片", use_column_width=True)

        # 预测按钮
        if st.button("识别品牌", type="primary"):
            with st.spinner("识别中..."):
                # 预处理和预测
                image_tensor = preprocess_image(image)
                probabilities = predict(model, device, image_tensor)

                # 获取前3个结果
                top3_probs, top3_classes = torch.topk(probabilities, 3)

                # 显示结果
                st.success("识别完成！")

                # 主要结果
                top_class = top3_classes[0][0].item()
                top_prob = top3_probs[0][0].item()

                st.metric(
                    label="预测品牌",
                    value=classnames.get(top_class, "未知"),
                    delta=f"{top_prob:.1%} 置信度"
                )

                # 其他可能结果
                st.subheader("其他可能品牌:")
                for i in range(1, 3):
                    class_idx = top3_classes[0][i].item()
                    prob = top3_probs[0][i].item()
                    st.write(f"{classnames.get(class_idx, '未知')}: {prob:.1%}")


if __name__ == "__main__":
    main()